module windump.registry;

import std.string;
import std.c.windows.windows;
import std.utf;

debug(WD_Unittest)
	import std.stdio;

//import windump.util;
import windump.utf;
import windump.common;

//WINAPI
extern(Windows)
{
    LONG RegOpenKeyExW(
            HKEY hKey,
            LPCWSTR lpSubKey,
            DWORD ulOptions,
            REGSAM samDesired,
            PHKEY phkResult
            );

    LONG RegCloseKey( HKEY hKey);

    LONG RegQueryValueExW( 
            HKEY hKey, 
            LPCWSTR lpValueName, 
            LPDWORD lpReserved, 
            LPDWORD lpType, 
            LPBYTE lpData, 
            LPDWORD lpcbData);

    LONG RegEnumKeyExW(
            HKEY hKey,
            DWORD dwIndex,
            LPWSTR lpName,
            LPDWORD lpcName,
            LPDWORD lpReserved,
            LPWSTR lpClass,
            LPDWORD lpcClass,
            PFILETIME lpftLastWriteTime
            );

    LONG RegEnumValueW(
            HKEY hKey,
            DWORD dwIndex,
            LPWSTR lpValueName,
            LPDWORD lpcValueName,
            LPDWORD lpReserved,
            LPDWORD lpType,
            LPBYTE lpData,
            LPDWORD lpcbData
            );

    LONG RegQueryInfoKeyW(
            HKEY hKey,
            LPWSTR lpClass,
            LPDWORD lpcClass,
            LPDWORD lpReserved,
            LPDWORD lpcSubKeys,
            LPDWORD lpcMaxSubKeyLen,
            LPDWORD lpcMaxClassLen,
            LPDWORD lpcValues,
            LPDWORD lpcMaxValueNameLen,
            LPDWORD lpcMaxValueLen,
            LPDWORD lpcbSecurityDescriptor,
            PFILETIME lpftLastWriteTime
            );

	LONG RegCreateKeyExW(
			HKEY hKey,
			LPCWSTR lpSubKey,
			DWORD Reserved,
			LPWSTR lpClass,
			DWORD dwOptions,
			REGSAM samDesired,
			LPSECURITY_ATTRIBUTES lpSecurityAttributes,
			PHKEY phkResult,
			LPDWORD lpdwDisposition
	);

	LONG RegSetValueExW(
			HKEY hKey,
			LPCWSTR lpValueName,
			DWORD Reserved,
			DWORD dwType,
			BYTE* lpData,
			DWORD cbData
	);

    // Registry APIs for D
}

class RegistryException : Exception
{
    public this(string msg)
    {
        super(msg);
    }
}

private HKEY RegOpenKey(HKEY hkey, string subkey, REGSAM sam)
{
    HKEY hret;
    wchar* name = null;

    LONG result = 
        RegOpenKeyExW(hkey, toUtf16z(subkey), 0, sam, &hret);
    if(result != ERROR_SUCCESS)
        throw new RegistryException("Failed to RegOpenKey");
    return hret;
}

private void RegClose(HKEY key)
{
   if(RegIsRootKey(key))
        return;

    RegCloseKey(key);
}

private Value RegQueryValue(HKEY hkey, string name)
{   
    DWORD type = 0, dataSize = 0;
    LONG r = RegQueryValueExW(hkey, toUtf16z(name), null, &type, null, &dataSize);
    if(r != ERROR_SUCCESS)
        throw new RegistryException("Failed to Query Value Size");

    ubyte[] data;
    data.length = dataSize;
    r = RegQueryValueExW(hkey, toUtf16z(name), null, &type, data.ptr, &dataSize);
    if(r != ERROR_SUCCESS)
        throw new RegistryException("Failed to Query Value Data");
    Value v = new Value(data, type);
    return v;
}

private string RegEnumKey(HKEY hkey, DWORD index)
{
    wchar[256] name;
    DWORD chars = name.length;
    LONG result = 
        RegEnumKeyExW(hkey, index, name.ptr, &chars, null, null, null, null);
    if(result != ERROR_SUCCESS)
        throw new RegistryException("Failed to RegEnumKey");
    string ret =  toUTF8(name[0..chars]);
    return ret;
}

private HKEY RegCreateKey(HKEY hkey, string name)
{
		HKEY hr = null;
		LONG result = 
		RegCreateKeyExW(hkey, toUtf16z(name), 0, null, 0, 
			KEY_ALL_ACCESS, null, &hr, null);
		if(result != ERROR_SUCCESS)
			throw new RegistryException("Failed to RegCreateKey");
		return hr;
}

private string RegEnumValue(HKEY hkey, DWORD index, ubyte[] data = null)
{
    wchar[16384] name;
    DWORD chars = name.length;
    DWORD type = 0;
    DWORD dataSize;
    
    LONG result = 
        RegEnumValueW(hkey, index, name.ptr, &chars, null, &type, null, &dataSize);
    if(result != ERROR_SUCCESS)
        throw new RegistryException("Failed to EnumValue");
    return toUTF8(name[0..chars]);
}

private void RegQueryKeyInfo(HKEY hkey, out uint subkeys, out uint values)
{
    LONG result = 
        RegQueryInfoKeyW(hkey, null, null, null, 
                &subkeys, null, null, &values, null, null, null, null);
    if(result != ERROR_SUCCESS)
        throw new RegistryException("Failed to RegQueryInfoKey");
}

private void RegSetValue(HKEY hkey, string name, ubyte[] data, DWORD type)
{
	LONG result = RegSetValueExW(hkey, toUtf16z(name), 0, type, data.ptr, data.length);
	if(result != ERROR_SUCCESS)
		throw new RegistryException("Failed to RegSetValue");
}

private uint RegValueCount(HKEY hkey)
{
    uint keys, values;
    RegQueryKeyInfo(hkey, keys, values);
    return values;
}

private uint RegKeyCount(HKEY hkey)
{
    uint keys, values;
    RegQueryKeyInfo(hkey, keys, values);
    return keys;
}

private bool RegIsRootKey(HKEY hkey)
{
    bool ret;
    switch(cast(DWORD)hkey)
    {
        case cast(DWORD)HKEY_CLASSES_ROOT:
        case cast(DWORD)HKEY_CURRENT_USER:
        case cast(DWORD)HKEY_LOCAL_MACHINE:
        case cast(DWORD)HKEY_USERS:
            ret = true;
            break;

        default:
            ret = false;
            break;
    }
    return ret;
}

class Value
{
    enum TypeIndex : int
    {
        DWord = 0,
        Binary,
        String,
    }
    private ubyte[] m_data;
    private DWORD   m_type;

    public this(ubyte[] data, DWORD type)
    {
        m_data = data;
        m_type = type;
    }

	public this(DWORD dwdata)
	{
		m_type = REG_DWORD;
		m_data.length = DWORD.sizeof;
		ubyte* ptr = cast(ubyte*)(&dwdata);
		m_data[] = ptr[0 .. m_data.length];
	}

	public this(ubyte[] data)
	{
		m_type = REG_BINARY;
		m_data = data;
	}

	public this(string data)
	{
		m_type = REG_SZ;
		wstring wstr = toUTF16(data);
		wstr ~= '\0';
		m_data.length = wstr.length * wchar.sizeof;
		m_data[] = (cast(ubyte*)wstr.ptr)[0 .. m_data.length];		
	}

	public DWORD type()
	{
		return m_type;
	}

	public ubyte[] data()
	{
		return m_data;
	}

    public string value_SZ()
    {
        assert(m_type == REG_SZ || m_type == REG_MULTI_SZ);
        wchar* ptr = cast(wchar*)m_data.ptr;

        version(D_Version2) { 
            //string ret = toUtf8(cast(immutable(wchar)*)ptr);
        }
        else string ret = toUtf8(ptr);

        return ret;
    }

    public ubyte[] value_BINARY()
    {
        assert(m_type == REG_BINARY);
        return m_data;
        //return m_var.get!(ubyte[])();
    }

    public DWORD value_DWORD()
    {
        //return m_var.get!(DWORD)();
        //FIXME
		assert(false);
        return 0;
    }
}


final class Registry
{
    private static Key m_hkcu = null;
    private static Key m_hklm = null;

    public static this() {
        m_hklm = new Key(HKEY_LOCAL_MACHINE, "HKEY_LOCAL_MACHINE");
        m_hkcu = new Key(HKEY_CURRENT_USER, "HKEY_CURRENT_USER");
    }

    private this() { }

    public static Key localMachine() {
        return m_hklm;
    }

    public static Key currentUser() {
        return m_hkcu;
    }

}

class Key : IDisposable
{
    private string m_name;
    private HKEY    m_handle = null;

    private this(HKEY parent, string name)
    in {
        assert(name !is null);
    }
    body {
        m_handle = parent;
        m_name = name;
    }

    public ~this()
    {
        if(m_handle !is null)
            close();
    }

    public void dispose()
    {
        close();
    }

    public HKEY handle() { 
        return m_handle;
    }

    public void close()
    {
        if(m_handle !is null) {
            RegClose(m_handle);
            m_handle = null;
        }
    }

    public string name() {
        return m_name;
    }
	
	public Key createSubkey(string name)
	{
		HKEY hkey = RegCreateKey(m_handle, name);
		Key subkey = new Key(hkey, name);
		return subkey;
	}

    public Key openSubkey(string name, bool writable = false)
    {
		DWORD flags = KEY_READ;
		if(writable) flags |= KEY_WRITE;
        HKEY hkey = RegOpenKey(m_handle, name, flags);
        Key subkey = new Key(hkey, name);
        return subkey;
    }

    public Key openSubkey(uint index, bool writable = false)
    {
        string name = RegEnumKey(m_handle, index);
        return openSubkey(name, writable);
    }

	public void setValue(string name, Value val)
	{
		RegSetValue(m_handle, name, val.data, val.type);
	}

	public void setValue(DWORD dwdata)
	{
		ubyte[DWORD.sizeof] data;
		ubyte* ptr = cast(ubyte*)(&dwdata);
		data[] = ptr[0 .. data.length];
		RegSetValue(m_handle, name, data, REG_DWORD);
	}

	public void setValue(string name, ubyte[] data)
	{
		RegSetValue(m_handle, name, data, REG_BINARY);
	}

	public void setValue(string name, string data)
	{
		wstring wstr = toUTF16(data);
		wstr ~= '\0';
		ubyte[] wdata;
		wdata.length = wstr.length * wchar.sizeof;
		wdata[] = (cast(ubyte*)wstr.ptr)[0 .. wdata.length];		
		RegSetValue(m_handle, name, wdata, REG_SZ);
	}

    public Value getValue(string name)
    {
        return RegQueryValue(m_handle, name);
    }

    public Value getValue(uint index)
    {
        string valueName = RegEnumValue(handle, index);
        return RegQueryValue(handle, valueName);
    }

    public string getValueName(uint index)
    {
        return RegEnumValue(m_handle, index);
    }

    public string[] subkeyNames()
    {
        string[] ret;
        for(uint i = 0; i < subkeyCount; i++)
        {
            ret ~= getSubkeyName(i);
        }

        return ret;
    }

    public string getSubkeyName(uint index)
    {
        return RegEnumKey(m_handle, index);
    }

    public string[] valueNames()
    {
        string[] ret;
        for(uint i = 0; i < valueCount; i++)
        {
            ret ~= getValueName(i);
        }
        return ret;
    }

    public uint valueCount()
    {
        return RegValueCount(m_handle);
    }

    public uint subkeyCount()
    {
        return RegKeyCount(m_handle);
    }

	public override string toString()
	{
		return name;
	}
}
